# -*- coding:utf8 -*-
# @Time : 2022/12/2 3:50 下午
# @Author : WanJie Wu

import jieba
import torch
from abc import abstractmethod


def load_stop_words(stop_words_path):
    if not stop_words_path:
        return []

    stop_words_lst = []
    with open(stop_words_path, "r", encoding="utf8") as f:
        for line in f.readlines():
            stop_words_lst.append(line.strip())
    return stop_words_lst


STOP_WORDS = load_stop_words("/note/nlp_algo/app/data/stop_words.txt")


class NLTK(object):

    @abstractmethod
    def ss_segment(self, *args, **kwargs):
        """单个句子分词 single sentence segment"""
        raise NotImplementedError


class SegmentJB(NLTK):
    """
    开源jieba库
    参考： https://github.com/fxsjy/jieba
    """

    def ss_segment(self, single_sentence, cut_all=False, hmm=True):
        """默认精确分词，采用HMM算法"""
        words = jieba.lcut(single_sentence, cut_all=cut_all, HMM=hmm)
        words = list(filter(lambda x: x not in STOP_WORDS, words))
        words = [word.strip() for word in words if len(word.strip()) >= 1]
        return words


class SegmentHLP(NLTK):
    """
    Han Language Processing
    参考： https://hanlp.hankcs.com/docs/api/hanlp/pretrained/index.html
    """
    def __init__(self, tok_model="COARSE_ELECTRA_SMALL_ZH"):
        import hanlp
        if tok_model:
            self._tok = hanlp.load(tok_model)

    def ss_segment(self, single_sentence):
        return self._tok(single_sentence)


class SegmentLTP(NLTK):
    """
    哈工大 Language Technology Platform
    参考： https://github.com/HIT-SCIR/ltp/blob/main/python/interface/README.md
    """
    def __init__(self, model_name="LTP/small"):
        from ltp import LTP
        self._ltp = LTP(model_name)
        if torch.cuda.is_available():
            self._ltp.to("cuda")

    def ss_segment(self, single_sentence):
        return self._ltp.pipeline(single_sentence, tasks=["cws"]).to_tuple()[0]


class SegmentLAC(NLTK):
    """百度词法分析工具, 全称Lexical Analysis of Chinese
    参考： https://github.com/baidu/lac
    """
    def __init__(self):
        from LAC import LAC
        self._seg = LAC(mode="seg")

    def ss_segment(self, single_sentence):
        return self._seg.run(single_sentence)


if __name__ == "__main__":
    texts = [
        "同学父亲去世我给五百，一月后我婆婆去世他的表现让我心寒",
        "古玩：地摊捡漏“八”大技巧",
        "司马光砸缸的故事大家都知道，司马光救的那个人是谁，你知道吗？",
        "宝宝出生三天没睁眼，奶奶用手掰开，我发现后制止，差点酿成大祸",
        "你碰到过的灵异事件，亲身经历的？"
    ]
    seg_obj = SegmentHLP()
    for text in texts:
        print(seg_obj.ss_segment(text))
